#!/usr/bin/env python3
"""
多波段功能测试脚本
测试多波段数据加载和模型输入功能
"""

import torch
import numpy as np
import yaml
import os
from data_utils import load_config, load_multi_band_himawari_data, find_himawari_file
from models import CombinedModel
from datetime import datetime

def test_multi_band_data_loading():
    """测试多波段数据加载功能"""
    print("=" * 50)
    print("测试多波段数据加载功能")
    print("=" * 50)
    
    # 加载配置
    config = load_config()
    band_list = config['data']['himawari_bands']
    print(f"配置中的波段列表: {band_list}")
    print(f"模型输入通道数: {config['model']['num_channels']}")
    
    # 查找一个测试文件
    test_time = datetime(2025, 6, 24, 0, 0)  # 使用示例时间
    himawari_path = find_himawari_file(test_time, config['data']['himawari_dir'], config['data'].get('output_nc_dir'))
    
    if himawari_path and os.path.exists(himawari_path):
        print(f"找到测试文件: {himawari_path}")
        
        # 测试多波段数据加载
        multi_band_data = load_multi_band_himawari_data(himawari_path, band_list)
        
        if multi_band_data is not None:
            print(f"多波段数据形状: {multi_band_data.shape}")
            print(f"波段数量: {multi_band_data.shape[0]}")
            print(f"空间分辨率: {multi_band_data.shape[1]} x {multi_band_data.shape[2]}")
            
            # 检查每个波段的数据范围
            for i, band in enumerate(band_list):
                band_data = multi_band_data[i]
                print(f"波段 {band}: 最小值={band_data.min():.3f}, 最大值={band_data.max():.3f}, 均值={band_data.mean():.3f}")
        else:
            print("多波段数据加载失败")
    else:
        print("未找到测试文件，创建模拟数据")
        # 创建模拟多波段数据
        num_bands = len(band_list)
        dummy_data = np.random.randn(num_bands, 224, 224).astype(np.float32)
        print(f"模拟多波段数据形状: {dummy_data.shape}")
        
    print()

def test_model_with_multi_band():
    """测试模型对多波段输入的处理"""
    print("=" * 50)
    print("测试模型多波段输入处理")
    print("=" * 50)
    
    # 加载配置
    config = load_config()
    
    # 创建模型
    model = CombinedModel(config)
    print(f"模型参数总数: {sum(p.numel() for p in model.parameters()):,}")
    
    # 创建测试数据
    batch_size = 2
    historical_points = config['data']['historical_days'] * 24 * 4  # 15分钟间隔
    num_channels = config['model']['num_channels']
    
    # 模拟场站数据 [batch_size, historical_points, 5]
    station_data = torch.randn(batch_size, historical_points, 5)
    
    # 模拟多波段卫星数据 [batch_size, channels, height, width]
    himawari_data = torch.randn(batch_size, num_channels, 224, 224)
    
    print(f"场站数据形状: {station_data.shape}")
    print(f"卫星数据形状: {himawari_data.shape}")
    
    # 测试前向传播
    try:
        with torch.no_grad():
            output = model(station_data, himawari_data)
            future_points = config['data']['future_hours'] * 4  # 15分钟间隔
            expected_shape = (batch_size, future_points, config['model']['num_output_vars'])
            
            print(f"模型输出形状: {output.shape}")
            print(f"期望输出形状: {expected_shape}")
            
            if output.shape == expected_shape:
                print("✅ 模型多波段输入测试通过")
            else:
                print("❌ 模型输出形状不匹配")
                
    except Exception as e:
        print(f"❌ 模型前向传播失败: {e}")
    
    print()

def test_data_consistency():
    """测试数据一致性"""
    print("=" * 50)
    print("测试数据一致性")
    print("=" * 50)
    
    config = load_config()
    
    # 检查配置一致性
    band_list = config['data']['himawari_bands']
    num_channels_config = config['model']['num_channels']
    
    print(f"配置波段数量: {len(band_list)}")
    print(f"模型输入通道数: {num_channels_config}")
    
    if len(band_list) == num_channels_config:
        print("✅ 配置一致性检查通过")
    else:
        print("❌ 配置不一致: 波段数量与模型输入通道数不匹配")
        print(f"  建议将模型num_channels设置为: {len(band_list)}")
    
    print()

def main():
    """主测试函数"""
    print("气象预测AI系统多波段功能测试")
    print("=" * 60)
    
    # 测试多波段数据加载
    test_multi_band_data_loading()
    
    # 测试数据一致性
    test_data_consistency()
    
    # 测试模型处理
    test_model_with_multi_band()
    
    print("=" * 60)
    print("测试完成！")
    print("如果所有测试都通过，可以开始多波段训练")

if __name__ == "__main__":
    main()
